#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/net.h>
#include <linux/socket.h>
#include <net/sock.h>
#include <net/tcp.h>
#include <linux/inet.h>  // For in_aton
#include <linux/mutex.h> // For mutex

#define MODBUS_TCP_DEVICE_NAME "modbus_tcp"

static int modbus_tcp_open(struct inode *inode, struct file *file);
static int modbus_tcp_release(struct inode *inode, struct file *file);
static ssize_t modbus_tcp_write(struct file *file, const char __user *buf, size_t count, loff_t *pos);
static ssize_t modbus_tcp_read(struct file *file, char __user *buf, size_t count, loff_t *pos);

static struct file_operations fops = {
    .owner = THIS_MODULE,
    .open = modbus_tcp_open,
    .release = modbus_tcp_release,
    .write = modbus_tcp_write,
    .read = modbus_tcp_read,
};

static struct cdev modbus_tcp_cdev;
static struct class *modbus_tcp_class;
static struct device *modbus_tcp_device;
static dev_t modbus_tcp_dev;
static DEFINE_MUTEX(modbus_tcp_mutex); // Mutex for synchronization

// Per-file private data structure
struct modbus_tcp_private {
    struct socket *socket;
};

static int modbus_tcp_open(struct inode *inode, struct file *file)
{
    struct modbus_tcp_private *priv;
    int ret;

    // Allocate private data
    priv = kzalloc(sizeof(*priv), GFP_KERNEL);
    if (!priv) {
        return -ENOMEM;
    }

    // Create a TCP socket
    ret = sock_create_kern(&init_net, AF_INET, SOCK_STREAM, IPPROTO_TCP, &priv->socket);
    if (ret < 0) {
        kfree(priv);
        printk(KERN_ERR "Failed to create socket\n");
        return ret;
    }

    // Connect to the Modbus TCP server
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(502);  // Modbus TCP default port
    server_addr.sin_addr.s_addr = in_aton("0.0.0.0");  // Replace with your server IP

    ret = kernel_connect(priv->socket, (struct sockaddr *)&server_addr, sizeof(server_addr), 0);
    if (ret < 0) {
        printk(KERN_ERR "Failed to connect to Modbus TCP server: %d\n", ret);
        sock_release(priv->socket);
        kfree(priv);
        return ret;
    }

    // Store private data in file->private_data
    file->private_data = priv;

    printk(KERN_INFO "Modbus TCP device opened\n");
    return 0;
}

static int modbus_tcp_release(struct inode *inode, struct file *file)
{
    struct modbus_tcp_private *priv = file->private_data;

    if (priv) {
        if (priv->socket) {
            kernel_sock_shutdown(priv->socket, SHUT_RDWR);
            sock_release(priv->socket);
        }
        kfree(priv);
    }

    printk(KERN_INFO "Modbus TCP device closed\n");
    return 0;
}

static ssize_t modbus_tcp_write(struct file *file, const char __user *buf, size_t count, loff_t *pos)
{
    struct modbus_tcp_private *priv = file->private_data;
    char buffer[1024];
    struct msghdr msg = {0};
    struct kvec iov = {buffer, min_t(size_t, count, sizeof(buffer))};
    struct iov_iter iter;
    int sent;

    if (!priv || !priv->socket) {
        printk(KERN_INFO "Modbus TCP socket not initialized\n");
        return -ENODEV;
    }

    if (copy_from_user(buffer, buf, min_t(size_t, count, sizeof(buffer)))) {
        return -EFAULT;
    }

    iov_iter_kvec(&iter, WRITE, &iov, 1, iov.iov_len);
    msg.msg_iter = iter;

    mutex_lock(&modbus_tcp_mutex); // Protect shared resources
    sent = sock_sendmsg(priv->socket, &msg);
    mutex_unlock(&modbus_tcp_mutex);

    if (sent < 0) {
        printk(KERN_INFO "Failed to send data, error: %d\n", sent);
        return sent;
    }

    printk(KERN_INFO "Data sent: %.*s\n", sent, buffer);
    //发送的数据打印一下，使用for循环
    int i = 0;
    for (i = 0; i < sent; i++) {
        printk(KERN_INFO "%02X ", buffer[i]);
    }
    printk(KERN_INFO "\n");
    return sent;
}

static ssize_t modbus_tcp_read(struct file *file, char __user *buf, size_t count, loff_t *pos)
{
    struct modbus_tcp_private *priv = file->private_data;
    char buffer[1024];
    struct msghdr msg = {0};
    struct kvec iov = {buffer, sizeof(buffer)};
    struct iov_iter iter;
    int received;

    if (!priv || !priv->socket) {
        printk(KERN_INFO "Modbus TCP socket not initialized\n");
        return -ENODEV;
    }

    iov_iter_kvec(&iter, READ, &iov, 1, iov.iov_len);
    msg.msg_iter = iter;

    mutex_lock(&modbus_tcp_mutex); // Protect shared resources
    received = sock_recvmsg(priv->socket, &msg, MSG_DONTWAIT);
    mutex_unlock(&modbus_tcp_mutex);

    if (received < 0) {
        printk(KERN_INFO "Failed to receive data, error: %d\n", received);
        return received;
    }

    if (copy_to_user(buf, buffer, min_t(size_t, received, count))) {
        return -EFAULT;
    }

    printk(KERN_INFO "Data received: %.*s\n", received, buffer);
    return received;
}

static int __init modbus_tcp_init(void)
{
    int ret;

    // Dynamically allocate a major number
    ret = alloc_chrdev_region(&modbus_tcp_dev, 0, 1, MODBUS_TCP_DEVICE_NAME);
    if (ret < 0) {
        printk(KERN_ERR "Failed to allocate character device region\n");
        return ret;
    }

    printk(KERN_INFO "Modbus TCP driver loaded with major number %d\n", MAJOR(modbus_tcp_dev));

    cdev_init(&modbus_tcp_cdev, &fops);
    modbus_tcp_cdev.owner = THIS_MODULE;

    ret = cdev_add(&modbus_tcp_cdev, modbus_tcp_dev, 1);
    if (ret < 0) {
        unregister_chrdev_region(modbus_tcp_dev, 1);
        printk(KERN_ERR "Failed to add character device\n");
        return ret;
    }

    modbus_tcp_class = class_create(THIS_MODULE, MODBUS_TCP_DEVICE_NAME);
    if (IS_ERR(modbus_tcp_class)) {
        cdev_del(&modbus_tcp_cdev);
        unregister_chrdev_region(modbus_tcp_dev, 1);
        printk(KERN_ERR "Failed to create class\n");
        return PTR_ERR(modbus_tcp_class);
    }

    modbus_tcp_device = device_create(modbus_tcp_class, NULL, modbus_tcp_dev, NULL, MODBUS_TCP_DEVICE_NAME);
    if (IS_ERR(modbus_tcp_device)) {
        class_destroy(modbus_tcp_class);
        cdev_del(&modbus_tcp_cdev);
        unregister_chrdev_region(modbus_tcp_dev, 1);
        printk(KERN_ERR "Failed to create device\n");
        return PTR_ERR(modbus_tcp_device);
    }

    printk(KERN_INFO "Modbus TCP driver initialized successfully\n");
    return 0;
}

static void __exit modbus_tcp_exit(void)
{
    device_destroy(modbus_tcp_class, modbus_tcp_dev);
    class_destroy(modbus_tcp_class);
    cdev_del(&modbus_tcp_cdev);
    unregister_chrdev_region(modbus_tcp_dev, 1);

    printk(KERN_INFO "Modbus TCP driver unloaded\n");
}

module_init(modbus_tcp_init);
module_exit(modbus_tcp_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Your Name");
MODULE_DESCRIPTION("A simple Modbus TCP driver");
MODULE_VERSION("0.1");